package com.gitee.wsl.flow.create

import androidx.annotation.RestrictTo
import com.gitee.wsl.flow.F
import kotlinx.coroutines.currentCoroutineContext
import kotlinx.coroutines.ensureActive
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.FlowCollector
import kotlinx.coroutines.flow.emitAll
import kotlinx.coroutines.flow.emptyFlow
import kotlinx.coroutines.flow.flow
import kotlin.time.Duration
import kotlinx.coroutines.delay as coroutinesDelay
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.flow.StateFlow
import kotlinx.coroutines.flow.drop
import kotlinx.coroutines.flow.onStart
import kotlinx.coroutines.flow.stateIn

/**
 * Creates a [Flow] that emits a sequence of numbers within a specified range.
 * @param start The value of the first integer in the sequence.
 * @param count The number of sequential integers to generate.
 */
 fun F.range(start: Int, count: Int): Flow<Int> =
    if (count <= 0) {
        emptyFlow()
    } else {
        flow { repeat(count) { emit(it + start) } }
    }

// ----------------------------------------------- REPEAT FOREVER -----------------------------------------------

/**
 * Returns a [Flow] that repeats all values emitted by the original [Flow] indefinitely.
 *
 * Note: If the source [Flow] is completed synchronously immediately (e.g. [emptyFlow]),
 * this will cause an infinite loop.
 */

 fun <T> Flow<T>.repeat(): Flow<T> =
    repeatInternal(
        flow = this,
        count = 0,
        infinite = true,
        delay = noDelay(),
    )

/**
 * Returns a [Flow] that repeats all values emitted by the original [Flow] indefinitely,
 * with a delay computed by [delay] function between each repetition.
 *
 * Note: If the source [Flow] is completed synchronously immediately (e.g. [emptyFlow]),
 * and [delay] returns [Duration.ZERO] or a negative value,
 * this will cause an infinite loop.
 */

 fun <T> Flow<T>.repeat(delay: suspend (count: Int) -> Duration): Flow<T> =
    repeatInternal(
        flow = this,
        count = 0,
        infinite = true,
        delay = delaySelector(delay),
    )

/**
 * Returns a [Flow] that repeats all values emitted by the original [Flow] indefinitely,
 * with a fixed [delay] between each repetition.
 *
 * Note: If the source [Flow] is completed synchronously immediately (e.g. [emptyFlow]),
 * and [delay] is [Duration.ZERO], this will cause an infinite loop.
 */

 fun <T> Flow<T>.repeat(delay: Duration): Flow<T> =
    repeatInternal(
        flow = this,
        count = 0,
        infinite = true,
        delay = fixedDelay(delay),
    )

// --------------------------------------------------- REPEAT COUNT ---------------------------------------------------

/**
 * Returns a [Flow] that repeats all values emitted by the original [Flow] at most [count] times.
 * If [count] is zero or negative, the resulting [Flow] completes immediately without emitting any items (i.e. [emptyFlow]).
 */

 fun <T> Flow<T>.repeat(count: Int): Flow<T> =
    repeatInternal(
        flow = this,
        count = count,
        infinite = false,
        delay = noDelay(),
    )

/**
 * Returns a [Flow] that repeats all values emitted by the original [Flow] at most [count] times,
 * with a delay computed by [delay] function between each repetition.
 *
 * If [count] is zero or negative, the resulting [Flow] completes immediately without emitting any items (i.e. [emptyFlow]).
 */

 fun <T> Flow<T>.repeat(
    count: Int,
    delay: suspend (count: Int) -> Duration,
): Flow<T> =
    repeatInternal(
        flow = this,
        count = count,
        infinite = false,
        delay = delaySelector(delay),
    )

/**
 * Returns a [Flow] that repeats all values emitted by the original [Flow] indefinitely,
 * with a fixed [delay] between each repetition.
 *
 * If [count] is zero or negative, the resulting [Flow] completes immediately without emitting any items (i.e. [emptyFlow]).
 */

 fun <T> Flow<T>.repeat(
    count: Int,
    delay: Duration,
): Flow<T> =
    repeatInternal(
        flow = this,
        count = count,
        infinite = false,
        delay = fixedDelay(delay),
    )

// ---------------------------------------------------- INTERNAL ----------------------------------------------------

private typealias DelayDurationSelector = suspend (count: Int) -> Duration

private inline fun noDelay(): DelayDurationSelector? = null

private inline fun fixedDelay(delay: Duration): DelayDurationSelector? =
    if (delay.isZeroOrNegative()) {
        noDelay()
    } else {
        FixedDelayDurationSelector(delay)
    }

private inline fun delaySelector(noinline delay: DelayDurationSelector): DelayDurationSelector = delay

private inline fun Duration.isZeroOrNegative() = this == Duration.ZERO || isNegative()

/**
 * Used when the delay duration is fixed.
 * This is an optimization to avoid integer overflow.
 */
private class FixedDelayDurationSelector(val duration: Duration) : DelayDurationSelector {
    override suspend fun invoke(count: Int): Duration = duration
}


private fun <T> repeatInternal(
    flow: Flow<T>,
    count: Int,
    infinite: Boolean,
    delay: DelayDurationSelector?,
): Flow<T> = when {
    infinite -> repeatIndefinitely(
        flow = flow,
        delay = delay,
    )
    count <= 0 -> emptyFlow()
    else -> repeatAtMostCount(
        flow = flow,
        count = count,
        delay = delay,
    )
}


private fun <T> repeatIndefinitely(
    flow: Flow<T>,
    delay: DelayDurationSelector?,
): Flow<T> = when (delay) {
    null -> flow {
        while (true) {
            emitAll(flow)
        }
    }
    is FixedDelayDurationSelector -> flow {
        while (true) {
            emitAll(flow)
            coroutinesDelay(delay.duration)
        }
    }
    else -> flow {
        var soFar = 1

        while (true) {
            emitAll(flow)
            coroutinesDelay(delay(soFar++))
        }
    }
}


private fun <T> repeatAtMostCount(
    flow: Flow<T>,
    count: Int,
    delay: DelayDurationSelector?,
): Flow<T> = when (delay) {
    null -> flow {
        repeat(count) {
            emitAll(flow)
        }
    }
    is FixedDelayDurationSelector -> flow {
        repeat(count) {
            emitAll(flow)
            coroutinesDelay(delay.duration)
        }
    }
    else -> flow {
        for (soFar in 1..count) {
            emitAll(flow)
            coroutinesDelay(delay(soFar))
        }
    }
}



/**
 * Creates a _cold_ flow that produces a single value from the given [function].
 * It calls [function] for each new [FlowCollector].
 *
 * This function is similar to [RxJava's fromCallable](http://reactivex.io/RxJava/3.x/javadoc/io/reactivex/rxjava3/core/Flowable.html#fromCallable-java.util.concurrent.Callable-).
 * See also [flowFromNonSuspend] for the non-suspend version.
 *
 * The returned [Flow] is cancellable and has the same behaviour as [kotlinx.coroutines.flow.cancellable].
 *
 * ## Example of usage:
 *
 * ```kotlin
 * suspend fun remoteCall(): R = ...
 * fun remoteCallFlow(): Flow<R> = flowFromSuspend(::remoteCall)
 * ```
 *
 * ## Another example:
 *
 * ```kotlin
 * var count = 0L
 * val flow = flowFromSuspend {
 *   delay(count)
 *   count++
 * }
 *
 * flow.collect { println("flowFromSuspend: $it") }
 * println("---")
 * flow.collect { println("flowFromSuspend: $it") }
 * println("---")
 * flow.collect { println("flowFromSuspend: $it") }
 * ```
 *
 * Output:
 *
 * ```none
 * flowFromSuspend: 0
 * ---
 * flowFromSuspend: 1
 * ---
 * flowFromSuspend: 2
 * ```
 *
 * @see flowFromNonSuspend
 */
 fun <T> F.suspendFlow(function: suspend () -> T): Flow<T> = FlowFromSuspend(function)

// We don't need to use `AbstractFlow` here because we only emit a single value without a context switch,
// and we guarantee all Flow's constraints: context preservation and exception transparency.
private class FlowFromSuspend<T>(private val function: suspend () -> T) : Flow<T> {
    override suspend fun collect(collector: FlowCollector<T>) {
        val value = function()
        currentCoroutineContext().ensureActive()
        collector.emit(value)
    }
}

 fun <T> F.flow(function: () -> T): Flow<T> = FlowFromNonSuspend(function)

// We don't need to use `AbstractFlow` here because we only emit a single value without a context switch,
// and we guarantee all Flow's constraints: context preservation and exception transparency.
private class FlowFromNonSuspend<T>(private val function: () -> T) : Flow<T> {
    override suspend fun collect(collector: FlowCollector<T>) {
        val value = function()
        currentCoroutineContext().ensureActive()
        collector.emit(value)
    }
}

/**
 * Creates a [Flow] that, on collection, calls a [Flow] factory to make a [Flow] for each new [FlowCollector].
 *
 * In some circumstances, waiting until the last minute (that is, until collection time)
 * to generate the [Flow] can ensure that collectors receive the freshest data.
 *
 * Example of usage:
 *
 * ```
 * suspend fun remoteCall1(): R1 = ...
 * suspend fun remoteCall2(r1: R1): R2 = ...
 *
 * fun example1(): Flow<R2> = defer {
 *   val r1 = remoteCall1()
 *   val r2 = remoteCall2(r1)
 *   flowOf(r2)
 * }
 *
 * fun example2(): Flow<R1> = defer { flowOf(remoteCall1()) }
 * ```
 */
fun <T> F.flow(flowFactory: suspend () -> Flow<T>): Flow<T> = flow { emitAll(flowFactory()) }


/**
 * Ignores all elements emitted by the source [Flow], only passes calls of `complete` or `error`.
 *
 * The returned flow does not emit any values. It completes normally when the source flow completes normally.
 * Otherwise, it completes with the same exception as the source flow.
 */
 fun <T> Flow<T>.ignoreElements(): Flow<Nothing> = IgnoreElementsFlow(this)

// We don't need to use `AbstractFlow` here because we don't emit any value.
private class IgnoreElementsFlow<T>(private val flow: Flow<T>) : Flow<Nothing> {
    override suspend fun collect(collector: FlowCollector<Nothing>) = flow.collect{}
}

class InitialValueFlow<T>(private val flow: Flow<T>) : Flow<T> by flow {
    fun dropInitialValue(): Flow<T> = drop(1)
    suspend fun asStateFlow(scope: CoroutineScope): StateFlow<T> = stateIn(scope)
}

@RestrictTo(RestrictTo.Scope.LIBRARY_GROUP)
fun <T> Flow<T>.asInitialValueFlow(value: T): InitialValueFlow<T> = InitialValueFlow(
    onStart { emit(value) },
)